package middleware

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"ehsm_kms/api"
	"encoding/base64"
	"errors"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/iancoleman/orderedmap"
	"go.uber.org/zap"
	"io"
	"sort"
)

func SignMiddleware(log *zap.Logger, apiKey string) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		action, _ := ctx.GetQuery("action")
		var reqBody []byte
		defer func() {
			log.Debug("sign middleware", zap.String("action", action), zap.ByteString("body", reqBody))
		}()

		if action == "" {
			api.HandleError(ctx, 400, errors.New("action empty"), nil)
			ctx.Abort()
			return
		}

		params := orderedmap.New()
		reqBody, _ = ctx.GetRawData()
		ctx.Request.Body = io.NopCloser(bytes.NewBuffer(reqBody))
		err := params.UnmarshalJSON(reqBody)
		if err != nil {
			api.HandleError(ctx, 400, errors.New("params error"), nil)
			ctx.Abort()
			return
		}

		var reqSign string
		if _sign, ok := params.Get("sign"); ok {
			if reqSign, ok = _sign.(string); ok {
				params.Delete("sign")
			}
		}

		signString := paramsSortStr(params)
		hmacSha256 := hmac.New(sha256.New, []byte(apiKey))
		hmacSha256.Write([]byte(signString))
		sign := base64.StdEncoding.EncodeToString(hmacSha256.Sum(nil))
		if reqSign != sign {
			api.HandleError(ctx, 400, errors.New("sign error"), nil)
			ctx.Abort()
			return
		}
		ctx.Next()
	}
}

func paramsSortStr(signParams *orderedmap.OrderedMap) string {
	var str string
	sortedSignParams := sortMap(signParams)
	for _, k := range sortedSignParams.Keys() {
		v, _ := sortedSignParams.Get(k)
		if k == "payload" {
			payload := v.(*orderedmap.OrderedMap)
			str += "&" + k + "=" + paramsSortStr(payload)
		} else {
			str += fmt.Sprintf("&%s=%v", k, v)
		}
	}
	if len(str) > 0 {
		str = str[1:] // Remove leading "&"
	}
	return str
}

func sortMap(oldmap *orderedmap.OrderedMap) *orderedmap.OrderedMap {
	newmap := orderedmap.New()
	keys := oldmap.Keys()
	sort.Strings(keys)
	for _, key := range keys {
		value, _ := oldmap.Get(key)
		newmap.Set(key, value)
	}
	return newmap
}
